package xyz.noark.codec.asm;

import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import xyz.noark.codec.output.CodedOutputStream;

import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

public class OutputStreamGenerator {
    private static final ASMClassLoader CLASS_LOADER = new ASMClassLoader();
    
    private static final Map<Class<?>, Class<?>> PRIMITIVE_WRAPPER_MAP = new HashMap<>();
    static {
        PRIMITIVE_WRAPPER_MAP.put(boolean.class, Boolean.class);
        PRIMITIVE_WRAPPER_MAP.put(byte.class, Byte.class);
        PRIMITIVE_WRAPPER_MAP.put(char.class, Character.class);
        PRIMITIVE_WRAPPER_MAP.put(short.class, Short.class);
        PRIMITIVE_WRAPPER_MAP.put(int.class, Integer.class);
        PRIMITIVE_WRAPPER_MAP.put(long.class, Long.class);
        PRIMITIVE_WRAPPER_MAP.put(float.class, Float.class);
        PRIMITIVE_WRAPPER_MAP.put(double.class, Double.class);
    }

    @SuppressWarnings("unchecked")
    public Class<? extends CodedOutputStream> generateOutputStream(Class<?> type) {
        if (Map.class.isAssignableFrom(type)) {
            return generateMapOutputStream(type);
        }
        if (Collection.class.isAssignableFrom(type)) {
            return generateCollectionOutputStream(type);
        }
        
        String className = "xyz/noark/codec/generated/Gen" + 
            Math.abs(type.getName().hashCode()) + "OutputStream";
        
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
        cw.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC, className, null, 
                Type.getInternalName(CodedOutputStream.class), null);
                
        generateConstructor(cw);
        generateEncodeMethod(cw, type);
        
        cw.visitEnd();
        
        byte[] code = cw.toByteArray();
        return (Class<? extends CodedOutputStream>) CLASS_LOADER.defineClass(
                className.replace('/', '.'), code);
    }
    
    @SuppressWarnings("unchecked")
    private Class<? extends CodedOutputStream> generateMapOutputStream(Class<?> type) {
        String className = "xyz/noark/codec/generated/MapOutputStream";
        
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
        cw.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC, className, null, 
                Type.getInternalName(CodedOutputStream.class), null);
                
        generateConstructor(cw);
        generateMapEncodeMethod(cw);
        
        cw.visitEnd();
        
        byte[] code = cw.toByteArray();
        return (Class<? extends CodedOutputStream>) CLASS_LOADER.defineClass(
                className.replace('/', '.'), code);
    }
    
    private void generateConstructor(ClassWriter cw) {
        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null);
        mv.visitCode();
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL, Type.getInternalName(CodedOutputStream.class), 
                "<init>", "()V", false);
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(1, 1);
        mv.visitEnd();
    }
    
    private void generateEncodeMethod(ClassWriter cw, Class<?> type) {
        String fieldGetterClassName = "xyz/noark/codec/generated/Gen" + 
            Math.abs(type.getName().hashCode()) + "FieldGetter";
        generateFieldGetterClass(type, fieldGetterClassName);
        
        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, "encode", 
                "(Ljava/lang/Object;)V", null, null);
        mv.visitCode();
        
        // 将参数转换为具体类型
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitTypeInsn(Opcodes.CHECKCAST, Type.getInternalName(type));
        mv.visitVarInsn(Opcodes.ASTORE, 2);
        
        // 创建FieldGetter实例
        mv.visitTypeInsn(Opcodes.NEW, fieldGetterClassName.replace('.', '/'));
        mv.visitInsn(Opcodes.DUP);
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL, fieldGetterClassName.replace('.', '/'), 
                "<init>", "()V", false);
        mv.visitVarInsn(Opcodes.ASTORE, 3);
                
        // 为每个字段生成序列化代码
        for (Field field : type.getDeclaredFields()) {
            if (!Modifier.isStatic(field.getModifiers())) {
                generateFieldEncoder(mv, field, fieldGetterClassName);
            }
        }
        
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(4, 4);
        mv.visitEnd();
    }
    
    private void generateFieldGetterClass(Class<?> type, String className) {
        String fieldGetterClassName = "xyz/noark/codec/generated/Gen" + 
            Math.abs(type.getName().hashCode()) + "FieldGetter";
        
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
        cw.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC, fieldGetterClassName, 
                null, "java/lang/Object", null);
        
        // 生成构造方法
        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null);
        mv.visitCode();
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL, "java/lang/Object", "<init>", "()V", false);
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(1, 1);
        mv.visitEnd();
        
        // 为每个字段生成getter方法
        for (Field field : type.getDeclaredFields()) {
            if (!Modifier.isStatic(field.getModifiers())) {
                generateFieldGetter(cw, field, type);
            }
        }
        
        cw.visitEnd();
        CLASS_LOADER.defineClass(fieldGetterClassName, cw.toByteArray());
    }
    
    private void generateFieldGetter(ClassWriter cw, Field field, Class<?> type) {
        String methodName = "get" + field.getName();
        String desc = "(L" + type.getName().replace('.', '/') + ";)" 
                + Type.getDescriptor(field.getType());
        
        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, methodName, desc, null, null);
        mv.visitCode();
        
        // 创建标签
        Label tryStart = new Label();
        Label tryEnd = new Label();
        Label catchHandler = new Label();
        Label endLabel = new Label();
        
        // 添加try-catch块
        mv.visitTryCatchBlock(tryStart, tryEnd, catchHandler, "java/lang/Exception");
        
        mv.visitLabel(tryStart);
        
        // 获取Field对象
        mv.visitLdcInsn(Type.getType(type));
        mv.visitLdcInsn(field.getName());
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Class", 
                "getDeclaredField", "(Ljava/lang/String;)Ljava/lang/reflect/Field;", false);
        mv.visitVarInsn(Opcodes.ASTORE, 2);
        
        // 设置accessible
        mv.visitVarInsn(Opcodes.ALOAD, 2);
        mv.visitInsn(Opcodes.ICONST_1);
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/reflect/Field", 
                "setAccessible", "(Z)V", false);
        
        // 获取字段值
        mv.visitVarInsn(Opcodes.ALOAD, 2);
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        
        if (field.getType().isPrimitive()) {
            String primitiveGetter = getPrimitiveGetter(field.getType());
            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/reflect/Field",
                    primitiveGetter, "(Ljava/lang/Object;)" + Type.getDescriptor(field.getType()), false);
        } else {
            mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/reflect/Field",
                    "get", "(Ljava/lang/Object;)Ljava/lang/Object;", false);
            mv.visitTypeInsn(Opcodes.CHECKCAST, Type.getInternalName(field.getType()));
        }
        
        mv.visitLabel(tryEnd);
        mv.visitJumpInsn(Opcodes.GOTO, endLabel);
        
        // 异常处理
        mv.visitLabel(catchHandler);
        mv.visitVarInsn(Opcodes.ASTORE, 3); // 存储异常
        
        // 抛出运行时异常
        mv.visitTypeInsn(Opcodes.NEW, "java/lang/RuntimeException");
        mv.visitInsn(Opcodes.DUP);
        mv.visitVarInsn(Opcodes.ALOAD, 3);
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL, "java/lang/RuntimeException", 
                "<init>", "(Ljava/lang/Throwable;)V", false);
        mv.visitInsn(Opcodes.ATHROW);
        
        mv.visitLabel(endLabel);
        
        // 返回值
        mv.visitInsn(getReturnOpcode(field.getType()));
        
        // 增加本地变量表大小以容纳异常处理
        mv.visitMaxs(4, 4);
        mv.visitEnd();
    }
    
    private String getPrimitiveGetter(Class<?> type) {
        if (type == boolean.class) return "getBoolean";
        if (type == byte.class) return "getByte";
        if (type == char.class) return "getChar";
        if (type == short.class) return "getShort";
        if (type == int.class) return "getInt";
        if (type == long.class) return "getLong";
        if (type == float.class) return "getFloat";
        if (type == double.class) return "getDouble";
        throw new IllegalArgumentException("Not a primitive type: " + type);
    }
    
    private void generateFieldEncoder(MethodVisitor mv, Field field, String fieldGetterClassName) {
        mv.visitVarInsn(Opcodes.ALOAD, 0); // this
        mv.visitVarInsn(Opcodes.ALOAD, 3); // FieldGetter实例
        mv.visitVarInsn(Opcodes.ALOAD, 2); // 对象实例
        
        // 使用正确的类名格式
        String className = fieldGetterClassName.replace('.', '/');
        
        // 调用getter方法
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, className,
                "get" + field.getName(), 
                "(L" + field.getDeclaringClass().getName().replace('.', '/') + ";)" 
                        + Type.getDescriptor(field.getType()),
                false);
        
        // 如果是对象类型，需要进行类型转换
        if (!field.getType().isPrimitive() && field.getType() != String.class) {
            mv.visitLdcInsn(Type.getType(field.getType()));
        }
        
        // 调用相应的写入方法
        String writeMethod = getWriteMethod(field.getType());
        String writeMethodDesc = getWriteMethodDescriptor(field.getType());
        
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, 
                Type.getInternalName(CodedOutputStream.class),
                writeMethod, writeMethodDesc, false);
    }
    
    private String getWriteMethod(Class<?> type) {
        if (type == boolean.class) return "writeBoolean";
        if (type == byte.class) return "writeByte";
        if (type == char.class) return "writeChar";
        if (type == short.class) return "writeShort";
        if (type == int.class) return "writeInt";
        if (type == long.class) return "writeLong";
        if (type == float.class) return "writeFloat";
        if (type == double.class) return "writeDouble";
        if (type == String.class) return "writeString";
        // 对于其他对象类型，使用writeObject方法
        return "writeObject";
    }
    
    private String getWriteMethodDescriptor(Class<?> type) {
        if (type.isPrimitive()) {
            return "(" + Type.getDescriptor(type) + ")V";
        }
        if (type == String.class) {
            return "(Ljava/lang/String;)V";
        }
        // 对于其他对象类型，需要传入对象和类型信息
        return "(Ljava/lang/Object;Ljava/lang/Class;)V";
    }
    
    private int getReturnOpcode(Class<?> type) {
        if (type == void.class) return Opcodes.RETURN;
        if (type == int.class) return Opcodes.IRETURN;
        if (type == long.class) return Opcodes.LRETURN;
        if (type == float.class) return Opcodes.FRETURN;
        if (type == double.class) return Opcodes.DRETURN;
        return Opcodes.ARETURN;
    }
    
    @SuppressWarnings("unchecked")
    private Class<? extends CodedOutputStream> generateCollectionOutputStream(Class<?> type) {
        String className = "xyz/noark/codec/generated/CollectionOutputStream";
        
        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
        cw.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC, className, null, 
                Type.getInternalName(CodedOutputStream.class), null);
                
        generateConstructor(cw);
        generateCollectionEncodeMethod(cw);
        
        cw.visitEnd();
        
        byte[] code = cw.toByteArray();
        return (Class<? extends CodedOutputStream>) CLASS_LOADER.defineClass(
                className.replace('/', '.'), code);
    }
    
    private void generateCollectionEncodeMethod(ClassWriter cw) {
        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, "encode", 
                "(Ljava/lang/Object;)V", null, null);
        mv.visitCode();
        
        // 将参数转换为Collection
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitTypeInsn(Opcodes.CHECKCAST, "java/util/Collection");
        mv.visitVarInsn(Opcodes.ASTORE, 2);
        
        // 写入Collection大小
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitVarInsn(Opcodes.ALOAD, 2);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Collection", 
                "size", "()I", true);
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, 
                Type.getInternalName(CodedOutputStream.class),
                "writeInt", "(I)V", false);
        
        // 获取迭代器
        mv.visitVarInsn(Opcodes.ALOAD, 2);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Collection", 
                "iterator", "()Ljava/util/Iterator;", true);
        mv.visitVarInsn(Opcodes.ASTORE, 3);
        
        // 开始循环
        Label loopStart = new Label();
        Label loopEnd = new Label();
        
        mv.visitLabel(loopStart);
        mv.visitVarInsn(Opcodes.ALOAD, 3);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Iterator", 
                "hasNext", "()Z", true);
        mv.visitJumpInsn(Opcodes.IFEQ, loopEnd);
        
        // 获取元素
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitVarInsn(Opcodes.ALOAD, 3);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Iterator", 
                "next", "()Ljava/lang/Object;", true);
        mv.visitLdcInsn(Type.getType(Object.class));
        
        // 写入元素
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, 
                Type.getInternalName(CodedOutputStream.class),
                "writeObject", "(Ljava/lang/Object;Ljava/lang/Class;)V", false);
        
        mv.visitJumpInsn(Opcodes.GOTO, loopStart);
        mv.visitLabel(loopEnd);
        
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(3, 4);
        mv.visitEnd();
    }

    private void generateMapEncodeMethod(ClassWriter cw) {
        MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, "encode", 
                "(Ljava/lang/Object;)V", null, null);
        mv.visitCode();
        
        // 将参数转换为Map
        mv.visitVarInsn(Opcodes.ALOAD, 1);
        mv.visitTypeInsn(Opcodes.CHECKCAST, "java/util/Map");
        mv.visitVarInsn(Opcodes.ASTORE, 2);
        
        // 写入Map大小
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitVarInsn(Opcodes.ALOAD, 2);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Map", 
                "size", "()I", true);
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, 
                Type.getInternalName(CodedOutputStream.class),
                "writeInt", "(I)V", false);
        
        // 获取entrySet迭代器
        mv.visitVarInsn(Opcodes.ALOAD, 2);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Map", 
                "entrySet", "()Ljava/util/Set;", true);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Set", 
                "iterator", "()Ljava/util/Iterator;", true);
        mv.visitVarInsn(Opcodes.ASTORE, 3);
        
        // 开始循环
        Label loopStart = new Label();
        Label loopEnd = new Label();
        
        mv.visitLabel(loopStart);
        mv.visitVarInsn(Opcodes.ALOAD, 3);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Iterator", 
                "hasNext", "()Z", true);
        mv.visitJumpInsn(Opcodes.IFEQ, loopEnd);
        
        // 获取entry
        mv.visitVarInsn(Opcodes.ALOAD, 3);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Iterator", 
                "next", "()Ljava/lang/Object;", true);
        mv.visitTypeInsn(Opcodes.CHECKCAST, "java/util/Map$Entry");
        mv.visitVarInsn(Opcodes.ASTORE, 4);
        
        // 写入key
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitVarInsn(Opcodes.ALOAD, 4);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Map$Entry", 
                "getKey", "()Ljava/lang/Object;", true);
        mv.visitLdcInsn(Type.getType(Object.class));
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, 
                Type.getInternalName(CodedOutputStream.class),
                "writeObject", "(Ljava/lang/Object;Ljava/lang/Class;)V", false);
        
        // 写入value
        mv.visitVarInsn(Opcodes.ALOAD, 0);
        mv.visitVarInsn(Opcodes.ALOAD, 4);
        mv.visitMethodInsn(Opcodes.INVOKEINTERFACE, "java/util/Map$Entry", 
                "getValue", "()Ljava/lang/Object;", true);
        mv.visitLdcInsn(Type.getType(Object.class));
        mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, 
                Type.getInternalName(CodedOutputStream.class),
                "writeObject", "(Ljava/lang/Object;Ljava/lang/Class;)V", false);
        
        mv.visitJumpInsn(Opcodes.GOTO, loopStart);
        mv.visitLabel(loopEnd);
        
        mv.visitInsn(Opcodes.RETURN);
        mv.visitMaxs(4, 5);
        mv.visitEnd();
    }

    private String generateValidClassName(Class<?> type) {
        // 生成一个简单的唯一标识符
        int hashCode = Math.abs(type.getName().hashCode());
        String uniqueId = String.format("%08x", hashCode);
        
        // 确保生成的类名是有效的Java标识符
        return "Gen" + uniqueId;
    }
} 